from strong_sort import StrongSORT
from strong_sort.utils.parser import get_config
from yolo.utils.torch_utils import select_device


class Sorter:

    def __init__(self, device, nr_sources=1, config_strongsort='', strong_sort_weights=''):
        self.config_strongsort = config_strongsort
        self.strong_sort_weights = strong_sort_weights
        self.outputs = None
        self.strongsort_list = []
        self.cfg = get_config()
        self.curr_frames, self.prev_frames = None, None
        device = select_device(device)
        self.nr_sources = nr_sources
        self.init_deep_sort(device=device)

    def init_deep_sort(self, nr_sources=1, device='', half=False):
        # initialize StrongSORT
        # cfg = get_config()
        # print(self.config_strongsort, "self.config_strongsort")
        self.cfg.merge_from_file(self.config_strongsort)
        # Create as many strong strong_sort instances as there are video sources
        for i in range(nr_sources):
            self.strongsort_list.append(
                StrongSORT(
                    self.strong_sort_weights,
                    device,
                    half,
                    max_dist=self.cfg.STRONGSORT.MAX_DIST,
                    max_iou_distance=self.cfg.STRONGSORT.MAX_IOU_DISTANCE,
                    max_age=self.cfg.STRONGSORT.MAX_AGE,
                    n_init=self.cfg.STRONGSORT.N_INIT,
                    nn_budget=self.cfg.STRONGSORT.NN_BUDGET,
                    mc_lambda=self.cfg.STRONGSORT.MC_LAMBDA,
                    ema_alpha=self.cfg.STRONGSORT.EMA_ALPHA,
                )
            )
            self.strongsort_list[i].model.warmup()
        self.outputs = [None] * nr_sources
        self.curr_frames, self.prev_frames = [None] * nr_sources, [None] * nr_sources
